import akshare as ak
import pandas as pd
import time
import os
from concurrent.futures import ThreadPoolExecutor
from multiprocessing import cpu_count
from sqlalchemy import create_engine


def create_mysql_engine():
    """
    创建数据库引擎对象
    :return: 新创建的数据库引擎对象
    """
    # 引擎参数信息
    host = 'localhost'
    user = 'root'
    passwd = 'XXyy%402023'
    port = '3306'
    db = 'shares'
 
    # 创建连接数据库stock的引擎对象
    db_engine = create_engine(
        'mysql+pymysql://{0}:{1}@{2}:{3}/{4}?charset=utf8'.format(user, passwd, host, port, db))
 
    # 返回引擎对象
    return db_engine

start_date="20200101"
today = time.strftime("%Y%m%d", time.localtime(time.time()))
engine = create_mysql_engine()

# 创建本地缓存目录
if not os.path.exists("cache"):
    os.mkdir("cache")

def update_one(code):
    cache_file = 'cache/stocks-%s.csv' % (code)
    print('load from %s' % (cache_file))
    if os.path.exists(cache_file):
        df = pd.read_csv(cache_file)
    else:
        df = ak.stock_zh_a_hist(symbol=str(code), period="daily", start_date=start_date, end_date=today, adjust="qfq")
        df.to_csv(cache_file)

    df.to_sql(name='stock_%s' % (code), con=engine, if_exists='replace', index=False)

    print('.', end='')

def update():
    # 首先加载所有的股票代码，大概7000支
    df = ak.stock_zh_a_spot_em()

    # 先按照名称和代码，排除不想要的
    filtered = []
    for index, row in df.iterrows():
        code = row['代码']
        name = row['名称']
        # 排除ST开头的和8开头的
        if int(code) >= 800000 or 'ST' in name:
            continue

        filtered.append(code)
    
    print('get %d stocks' % (len(filtered)))

    pool = ThreadPoolExecutor(max_workers=cpu_count())

    for item in filtered:
        pool.submit(update_one, item)

    pool.shutdown(wait=True)

    print('Done')

if __name__=='__main__':
    update()

